{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import transforms, utils\n",
    "\n",
    "import time\n",
    "import numpy as np\n",
    "import os\n",
    "import pickle\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "plt.ion()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dataset_dir = './Dataset/'\n",
    "batch_size = 4\n",
    "epochs = 10\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "# device = 'cpu'\n",
    "learning_rate = 0.01"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class YouTubePose(Dataset):\n",
    "    \n",
    "    def __init__(self, datapoint_pairs, shapeLoss_datapoint_pairs, dataset_dir, transform=None, mode='train'):\n",
    "        self.datapoint_pairs = datapoint_pairs\n",
    "        self.shapeLoss_datapoint_pairs = shapeLoss_datapoint_pairs\n",
    "        self.dataset_dir = dataset_dir\n",
    "        self.transform = transform\n",
    "        self.mode = mode\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.datapoint_pairs)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        image_pair = self.datapoint_pairs[idx]\n",
    "        x_gen_path = image_pair[0]\n",
    "        x_dis_path = image_pair[1]\n",
    "        y_path = image_pair[2]\n",
    "        \n",
    "        identity_pair = self.shapeLoss_datapoint_pairs[idx]\n",
    "        iden_1_path = identity_pair[0]\n",
    "        iden_2_path = identity_pair[1]\n",
    "        \n",
    "        x_gen = Image.open(self.dataset_dir + self.mode + '/' + x_gen_path)\n",
    "        x_dis = Image.open(self.dataset_dir + self.mode + '/' + x_dis_path)\n",
    "        y = Image.open(self.dataset_dir + self.mode + '/' + y_path)\n",
    "        iden_1 = Image.open(self.dataset_dir + self.mode + '/' + iden_1_path)\n",
    "        iden_2 = Image.open(self.dataset_dir + self.mode + '/' + iden_2_path)\n",
    "        \n",
    "        if self.transform:\n",
    "            x_gen = self.transform(x_gen)\n",
    "            x_dis = self.transform(x_dis)\n",
    "            y = self.transform(y)\n",
    "            iden_1 = self.transform(iden_1)\n",
    "            iden_2 = self.transform(iden_2)\n",
    "            \n",
    "        sample = {'x_gen' : x_gen, 'x_dis': x_dis, 'y': y, 'iden_1': iden_1, 'iden_2':iden_2}\n",
    "        return sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./dataset_lists/train_datapoint_triplets.pkl', 'rb') as f:\n",
    "    datapoint_pairs = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./dataset_lists/train_shapeLoss_pairs.pkl', 'rb') as f:\n",
    "    shapeLoss_datapoint_pairs = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(datapoint_pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.Resize((256, 256)),\n",
    "    transforms.ToTensor()\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_dataset = YouTubePose(datapoint_pairs, shapeLoss_datapoint_pairs, dataset_dir, transform)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True,\n",
    "                                             num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dataset_sizes = [len(train_dataset)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_landmarks_batch(sample_batched):\n",
    "    \n",
    "    images_batch, landmarks_batch = sample_batched['iden_1'], sample_batched['iden_2']\n",
    "#     batch_size = len(images_batch)\n",
    "#     im_size = images_batch.size(2)\n",
    "    \n",
    "    grid = utils.make_grid(images_batch)\n",
    "    grid1 = utils.make_grid(landmarks_batch)\n",
    "    \n",
    "#     plt.imshow(grid.numpy().transpose((1, 2, 0)))\n",
    "    plt.imshow(grid1.numpy().transpose((1, 2, 0)))\n",
    "    plt.title('Batch from dataloader')\n",
    "    \n",
    "for i_batch, sample_batched in enumerate(train_dataloader):\n",
    "    print(i_batch, sample_batched['x_gen'].size(), sample_batched['y'].size())\n",
    "    \n",
    "    if i_batch == 3:\n",
    "        plt.figure()\n",
    "        show_landmarks_batch(sample_batched)\n",
    "        plt.axis('off')\n",
    "        plt.ioff()\n",
    "        plt.show()\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 3x3 convolution\n",
    "def conv3x3(in_channels, out_channels, stride=1):\n",
    "    return nn.Conv2d(in_channels, out_channels, kernel_size=3, \n",
    "                     stride=stride, padding=1, bias=False)\n",
    "\n",
    "# Residual block\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, stride=1, downsample=None):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv1 = conv3x3(in_channels, out_channels, stride)\n",
    "        self.bn1 = nn.BatchNorm2d(out_channels)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv2 = conv3x3(out_channels, out_channels)\n",
    "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
    "        self.downsample = downsample\n",
    "        \n",
    "    def forward(self, x):\n",
    "#         if torch.cuda.is_available():\n",
    "#             x = torch.cuda.FloatTensor(x)\n",
    "        residual = x\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "        if self.downsample:\n",
    "            residual = self.downsample(x)\n",
    "        out += residual\n",
    "        out = self.relu(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self, block):\n",
    "        super(Generator, self).__init__()\n",
    "        \n",
    "        self.conv1_x = nn.Sequential(\n",
    "            nn.Conv2d(3, 16, 3, padding = 1),\n",
    "            nn.BatchNorm2d(16),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.AvgPool2d(2))\n",
    "        \n",
    "        self.conv2_x = nn.Sequential(\n",
    "            nn.Conv2d(16, 32, 3, padding = 1),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.AvgPool2d(2))\n",
    "\n",
    "        self.conv3_x = nn.Sequential(\n",
    "            nn.Conv2d(32, 16, 3, padding = 1),\n",
    "            nn.BatchNorm2d(16),\n",
    "            nn.ReLU(inplace=True))\n",
    "\n",
    "        self.conv1_y = nn.Sequential(\n",
    "            nn.Conv2d(3, 16, 3, padding = 1),\n",
    "            nn.BatchNorm2d(16),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.AvgPool2d(2))\n",
    "        \n",
    "        self.conv2_y = nn.Sequential(\n",
    "            nn.Conv2d(16, 32, 3, padding = 1),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.AvgPool2d(2))\n",
    "\n",
    "        self.conv3_y = nn.Sequential(\n",
    "            nn.Conv2d(32, 16, 3, padding = 1),\n",
    "            nn.BatchNorm2d(16),\n",
    "            nn.ReLU(inplace=True))      \n",
    "        \n",
    "#         self.conv1_x = nn.Conv2d(3, 16, 3, padding = 1)\n",
    "#         self.conv2_x = nn.Conv2d(16, 32, 3)\n",
    "#         self.conv3_x = nn.Conv2d(32, 16, 3, padding=1)\n",
    "\n",
    "#         self.conv1_y = nn.Conv2d(3, 16, 3)\n",
    "#         self.conv2_y = nn.Conv2d(16, 32, 3)\n",
    "#         self.conv3_y = nn.Conv2d(32, 16, 3, padding = 1)\n",
    "\n",
    "#         self.relu = nn.ReLU()\n",
    "#         self.avgpool = nn.AvgPool2d(2)\n",
    "        \n",
    "    \n",
    "        # 2 Residual Blocks for Identity Image\n",
    "        self.block1_x = block(16, 16)\n",
    "#         downsample_x = nn.Sequential(conv3x3(16, 1, 1), nn.BatchNorm2d(1))\n",
    "#         self.block2_x = block(16, 1, 1, downsample_x)\n",
    "        self.block2_x = block(16, 16)\n",
    "\n",
    "        # 2 Residual Blocks for Shape Image\n",
    "        self.block1_y = block(16, 16)\n",
    "#         downsample_y = nn.Sequential(conv3x3(16, 1, 1), nn.BatchNorm2d(1))\n",
    "#         self.block2_y = block(16, 1, 1, downsample_y)\n",
    "        self.block2_y = block(16, 16)\n",
    "        # 2 Residual Blocks for Combined(concat) image\n",
    "        downsample1_concat = nn.Sequential(conv3x3(32, 16, 1), nn.BatchNorm2d(16))\n",
    "        self.block1_concat = block(32, 16, 1, downsample1_concat)\n",
    "\n",
    "        self.block2_concat = block(16, 16)\n",
    "        \n",
    "#         self.deconv1 = nn.ConvTranspose2d(16, 16, 3)\n",
    "#         self.deconv2 = nn.ConvTranspose2d(16, 3, 3)\n",
    "        \n",
    "    \n",
    "#          self.conv2_y = nn.Sequential(\n",
    "#             nn.Conv2d(16, 32, 3, padding = 1),\n",
    "#             nn.BatchNorm2d(32),\n",
    "#             nn.ReLU(inplace=True),\n",
    "#             nn.AvgPool2d(2))\n",
    "    \n",
    "        self.upsample1 = nn.Sequential(\n",
    "            nn.Upsample(scale_factor=2, mode='bilinear', align_corners = True),\n",
    "            nn.ConvTranspose2d(16, 32, 3, padding=1),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.ReLU(inplace=True))\n",
    "        \n",
    "        self.upsample2 = nn.Sequential(\n",
    "            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),\n",
    "            nn.ConvTranspose2d(32, 3, 3, padding=1),\n",
    "            nn.BatchNorm2d(3),\n",
    "            nn.ReLU(inplace=True))\n",
    "        \n",
    "        \n",
    "    def forward(self, x, y):\n",
    "        \n",
    "        x = self.conv1_x(x)\n",
    "        x = self.conv2_x(x)\n",
    "        x = self.conv3_x(x)\n",
    "        x = self.block1_x(x)\n",
    "        x = self.block2_x(x)\n",
    "        \n",
    "        y = self.conv1_y(y)\n",
    "        y = self.conv2_y(y)\n",
    "        y = self.conv3_y(y)\n",
    "        y = self.block1_y(y)\n",
    "        y = self.block2_y(y)\n",
    "        \n",
    "#         if torch.cuda.is_available():\n",
    "#             concat_result = torch.cuda.FloatTensor([x.shape[0], x.shape[1] * 2, x.shape[2], x.shape[3]]).fill_(0)\n",
    "#         else:\n",
    "        concat_result = torch.zeros([x.shape[0], x.shape[1] * 2, x.shape[2], x.shape[3]], dtype=x.dtype)\n",
    "#         print(x.shape, y.shape, concat_result.shape)\n",
    "        for i in range(batch_size):\n",
    "            for j in range(x.shape[1]):\n",
    "                concat_result[i][j] = x[i][j]\n",
    "                concat_result[i][j + x.shape[1]] = y[i][j]\n",
    "        if torch.cuda.is_available():\n",
    "            concat_result = concat_result.cuda()\n",
    "        concat_result = self.block1_concat(concat_result)\n",
    "        concat_result = self.block2_concat(concat_result)\n",
    "        \n",
    "        upsampled_1 = self.upsample1(concat_result)\n",
    "        upsampled_2 = self.upsample2(upsampled_1)\n",
    "#         print(upsample2.shape)\n",
    "        return upsampled_2\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NLayerDiscriminator(nn.Module):\n",
    "    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False):\n",
    "        super(NLayerDiscriminator, self).__init__()\n",
    "        \n",
    "        use_bias = norm_layer\n",
    "\n",
    "        kw = 3\n",
    "        padw = 1\n",
    "        sequence = []\n",
    "#         sequence = [\n",
    "#             nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),\n",
    "#             nn.LeakyReLU(0.2, True)\n",
    "#         ]\n",
    "        self.initial_conv = nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw)\n",
    "        self.initial_relu = nn.LeakyReLU(0.2, True)\n",
    "        \n",
    "    \n",
    "        ndf = 128\n",
    "        \n",
    "        nf_mult = 1\n",
    "        nf_mult_prev = 1\n",
    "        for n in range(1, n_layers):\n",
    "            nf_mult_prev = nf_mult\n",
    "            nf_mult = min(2**n, 8)\n",
    "            sequence += [\n",
    "                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,\n",
    "                          kernel_size=kw, stride=2, padding=padw, bias=use_bias),\n",
    "                norm_layer(ndf * nf_mult),\n",
    "                nn.LeakyReLU(0.2, True)\n",
    "            ]\n",
    "\n",
    "        nf_mult_prev = nf_mult\n",
    "        nf_mult = min(2**n_layers, 8)\n",
    "        sequence += [\n",
    "            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,\n",
    "                      kernel_size=kw, stride=1, padding=padw, bias=use_bias),\n",
    "            norm_layer(ndf * nf_mult),\n",
    "            nn.LeakyReLU(0.2, True)\n",
    "        ]\n",
    "\n",
    "        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]\n",
    "\n",
    "        if use_sigmoid:\n",
    "            sequence += [nn.Sigmoid()]\n",
    "\n",
    "        self.model = nn.Sequential(*sequence)\n",
    "        \n",
    "        self.maxpool = nn.MaxPool2d(8)\n",
    "    def forward(self, x, y):\n",
    "        x = self.initial_relu(self.initial_conv(x))\n",
    "        y = self.initial_relu(self.initial_conv(y)) \n",
    "\n",
    "#         if torch.cuda.is_available():\n",
    "#             concat_fmap = torch.cuda.FloatTensor([x.shape[0], 2 * x.shape[1], x.shape[2], x.shape[3]]).fill_(0)\n",
    "#         else:\n",
    "        concat_fmap = torch.zeros([x.shape[0], 2 * x.shape[1], x.shape[2], x.shape[3]], dtype=x.dtype)\n",
    "\n",
    "            #         print(x.shape, y.shape, concat_result.shape)\n",
    "        for i in range(x.shape[0]):\n",
    "            for j in range(x.shape[1]):\n",
    "                concat_fmap[i][j] = x[i][j]\n",
    "                concat_fmap[i][j + 64] = y[i][j]\n",
    "        if torch.cuda.is_available():\n",
    "            concat_fmap = concat_fmap.cuda()\n",
    "        op = self.model(concat_fmap)\n",
    "        op_neg = - op\n",
    "        op_neg = self.maxpool(op_neg)\n",
    "        op_pooled = -op_neg\n",
    "        \n",
    "        return op, op_pooled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "generator = Generator(ResidualBlock)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "discriminator = NLayerDiscriminator(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# sample = []\n",
    "# for i_batch, sample_batched in enumerate(train_dataloader):\n",
    "#     sample = sample_batched\n",
    "#     if i_batch == 0:\n",
    "#         break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# op = generator(sample['x_gen'], sample['y'])\n",
    "# op.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "generator = generator.to(device)\n",
    "discriminator = discriminator.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# def lossIdentity(real_pooled_op, fake_pooled_op):\n",
    "#     batch_size = real_pooled_op.size()[0]\n",
    "#     real_pooled_op = 1 - real_pooled_op\n",
    "#     real_pooled_op = real_pooled_op ** 2\n",
    "#     fake_pooled_op  = (fake_pooled_op ** 2)\n",
    "#     real_pooled_op = real_pooled_op.view(batch_size, -1)\n",
    "#     fake_pooled_op = fake_pooled_op.view(batch_size, -1)\n",
    "#     real_pooled_op = torch.sum(real_pooled_op, dim = 1)\n",
    "#     fake_pooled_op = torch.sum(fake_pooled_op, dim = 1)\n",
    "#     return (real_pooled_op + fake_pooled_op) / batch_size\n",
    "\n",
    "def lossIdentity(real_pair, fake_pair):\n",
    "    batch_size = real_pair.size()[0]\n",
    "    real_pair = 1 - real_pair\n",
    "    real_pair = real_pair ** 2\n",
    "    fake_pair  = fake_pair ** 2\n",
    "    real_pair = torch.sum(real_pair)\n",
    "    fake_pair = torch.sum(fake_pair)\n",
    "    return (real_pair + fake_pair) / batch_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def lossShape(x, y):\n",
    "    batch_size = x.size()[0]\n",
    "    diff = x - y\n",
    "    diff = diff ** 2\n",
    "    diff = torch.sum(diff) / batch_size\n",
    "    return diff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(gen, disc, loss_i, loss_s, optimizer_gen, optimizer_disc, alpha = 1, beta = 1, num_epochs = 10):\n",
    "    for epoch in range(num_epochs):\n",
    "        print(\"Epoch {}/{}\".format(epoch, num_epochs - 1))\n",
    "        print('-'*10)\n",
    "        dataloader = train_dataloader\n",
    "        gen.train()\n",
    "        disc.train()\n",
    "        since = time.time()\n",
    "        running_loss_iden = 0.0\n",
    "        running_loss_s1 = 0.0\n",
    "        running_loss_s2a = 0.0\n",
    "        running_loss_s2b = 0.0\n",
    "        running_loss = 0.0\n",
    "        \n",
    "        for i_batch, sample_batched in enumerate(dataloader):\n",
    "            x_gen, y, x_dis = sample_batched['x_gen'], sample_batched['y'], sample_batched['x_dis']\n",
    "            iden_1, iden_2 = sample_batched['iden_1'], sample_batched['iden_2']\n",
    "            x_gen = x_gen.to(device)\n",
    "            y = y.to(device)\n",
    "            x_dis = x_dis.to(device)\n",
    "            iden_1 = iden_1.to(device)\n",
    "            iden_2 = iden_2.to(device)\n",
    "            \n",
    "            optimizer_gen.zero_grad()\n",
    "            optimizer_disc.zero_grad()\n",
    "            \n",
    "            with torch.set_grad_enabled(True):\n",
    "                x_generated = gen(x_gen, y)\n",
    "                print('forward 1 done')\n",
    "                fake_op, fake_pooled_op = disc(x_gen, x_generated)\n",
    "                real_op, real_pooled_op = disc(x_gen, x_dis)\n",
    "                loss_identity_gen = -loss_i(real_pooled_op, fake_pooled_op)\n",
    "                print('Loss calculated')\n",
    "                loss_identity_gen.backward(retain_graph=True)\n",
    "                optimizer_gen.step()\n",
    "                print('backward 1.1 done')\n",
    "                \n",
    "                optimizer_disc.zero_grad()\n",
    "                loss_identity_disc = loss_i(real_op, fake_op)\n",
    "                print('Loss calculated')\n",
    "                loss_identity_disc.backward(retain_graph=True)\n",
    "                optimizer_disc.step()\n",
    "                print('backward 1.2 done')\n",
    "\n",
    "                optimizer_gen.zero_grad()\n",
    "                optimizer_disc.zero_grad()\n",
    "                x_ls2a = gen(y, x_generated)\n",
    "                x_ls2b = gen(x_generated, y)\n",
    "                print('forward 2 done')\n",
    "\n",
    "                loss_s2a = loss_s(y, x_ls2a)\n",
    "                loss_s2b = loss_s(x_generated, x_ls2b)\n",
    "                loss_s2 = loss_s2a + loss_s2b\n",
    "                print('Loss calculated')\n",
    "\n",
    "                loss_s2.backward()\n",
    "                optimizer_gen.step()\n",
    "                print('backward 2 done')\n",
    "\n",
    "                optimizer_gen.zero_grad()\n",
    "                optimizer_disc.zero_grad()\n",
    "                \n",
    "                x_ls1 = generator(iden_1, iden_2)\n",
    "                print('forward 3 done')\n",
    "\n",
    "                loss_s1 = loss_s(iden_2, x_ls1)\n",
    "                print('Loss calculated')\n",
    "                loss_s1.backward()\n",
    "                optimizer_gen.step()\n",
    "                print('backward 5 done')\n",
    "                print()\n",
    "            running_loss_iden += loss_identity_disc.item() * x_gen.size(0)\n",
    "            running_loss_s1 += loss_s1.item() * x_gen.size(0)\n",
    "            running_loss_s2a += loss_s2a.item() * x_gen.size(0) \n",
    "            running_loss_s2b += loss_s2b.item() * x_gen.size(0)\n",
    "            running_loss = running_loss_iden +  beta * (running_loss_s1 + alpha * (running_loss_s2a + running_loss_s2b))\n",
    "            print(str(time.time() - since))\n",
    "            since = time.time()\n",
    "        epoch_loss_iden = running_loss_iden / dataset_sizes[0]\n",
    "        epoch_loss_s1 = running_loss_s1 / dataset_sizes[0]\n",
    "        epoch_loss_s2a = running_loss_s2a / dataset_sizes[0]\n",
    "        epoch_loss_s2b = running_loss_s2a / dataset_sizes[0]\n",
    "        epoch_loss = running_loss / dataset_sizes[0]\n",
    "        print('Identity Loss: {:.4f} Loss Shape1: {:.4f} Loss Shape2a: {:.4f} \\\n",
    "               Loss Shape2b: {:.4f}'.format(epoch_loss_iden, epoch_loss_s1,\n",
    "                                           epoch_loss_s2a, epoch_loss_s2b))\n",
    "        print('Epoch Loss: {:.4f}'.format(epoch_loss))\n",
    "        print('Time Taken: ' + str(time.time() - since))\n",
    "    return gen, disc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer_gen = optim.SGD(generator.parameters(), lr = learning_rate, momentum=0.9)\n",
    "optimizer_disc = optim.SGD(discriminator.parameters(), lr = learning_rate, momentum=0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator, discriminator = train_model(generator, discriminator, lossIdentity, lossShape, optimizer_gen, optimizer_disc, num_epochs=epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
